CassandraClientProvider.java

package net.ameba.cassandra.web.service;

import java.io.IOException;
import java.net.ConnectException;
import java.util.HashMap;
import java.util.Map;

import org.apache.cassandra.thrift.Cassandra.Client;
import org.apache.cassandra.tools.NodeProbe;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TFramedTransport;
import org.apache.thrift.transport.TSocket;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.TTransportException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

/**
 * {@link CassandraClientProvider} provides {@link Client} and {@link NodeProbe}.
 *
 * @author suguru
 */
@Component
public class CassandraClientProvider {
	
	@Autowired
	private CassandraProperties properties;
	
	private ThreadLocal<Client> th = new ThreadLocal<Client>();

	private Map<String, NodeProbe> probeMap;
	
	public CassandraClientProvider() {
		probeMap = new HashMap<String, NodeProbe>();
	}
	
	/**
	 * Returns true if client has been already created.
	 * @return
	 */
	public boolean hasThriftClient() {
		return th.get() != null;
	}
	
	/**
	 * Cleaning opened thrift client
	 */
	public void clean() {
		th.set(null);
	}

	/**
	 * Get Thrift Client
	 * @return
	 * @throws TTransportException
	 * @throws ConnectException
	 */
	public Client getThriftClient() throws TTransportException, ConnectException {
		
		if (!properties.hasProperties()) {
			return null;
		}
		
		String host = properties.getProperty(CassandraProperties.HOST);
		int port = Integer.parseInt(properties.getProperty(CassandraProperties.THRIFT_PORT));
		
		// Check usage of framed transport
		String framed = properties.getProperty(CassandraProperties.FRAMED_TRANSPORT);
		boolean isFramed = (framed != null && Boolean.valueOf(framed));
		
		Client client = th.get();
		if (client == null) {
			TTransport transport = new TSocket(host, port);
			((TSocket) transport).setTimeout(10000);
			// wrap transport if framed transport is enabled
			if (isFramed) {
				transport = new TFramedTransport(transport);
			}
			TProtocol protocol = new TBinaryProtocol(transport);
			transport.open();
			client = new Client(protocol);
			th.set(client);
		}
		return client;
	}
	
	/**
	 * Get NodeProbe connected to default host and port.
	 * @return
	 */
	public NodeProbe getProbe() {
		String host = properties.getProperty("cassandra.host");
		return getProbe(host);
	}
	
	/**
	 * Get NodeProbed connected to specified host.
	 * 
	 * @param host
	 * @return
	 */
	public synchronized NodeProbe getProbe(String host) {
		
		int port = Integer.parseInt(properties.getProperty("cassandra.jmx.port"));
		
		NodeProbe probe = probeMap.get(host + ":" + port);
		if (probe == null) {
			try {
				probe = new NodeProbe(host, port);
				probeMap.put(host + ":" + port, probe);
			} catch (IOException e) {
				return null;
			} catch (Exception e) {
				throw new RuntimeException(e);
			}
		} else {
			try {
				// check the connectivity
				probe.getToken();
			} catch (RuntimeException e) { 
				// clean jmx if failed to get token
				probeMap.remove(host + ":" + port);
				// retry
				try {
					probe = new NodeProbe(host, port);
					probeMap.put(host + ":" + port, probe);
				} catch (Exception ee) {
					return null;
				}
			}
		}
		return probe;
	}

}